package com.jt.dify.config;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import org.noear.solon.ai.chat.ChatChoice;
import org.noear.solon.ai.chat.ChatConfig;
import org.noear.solon.ai.chat.ChatOptions;
import org.noear.solon.ai.chat.ChatResponseDefault;
import org.noear.solon.ai.chat.dialect.AbstractChatDialect;
import org.noear.solon.ai.chat.message.AssistantMessage;
import org.noear.solon.ai.chat.message.ChatMessage;

import java.util.Date;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

/**
 * @author by HaiTao.Wang on 2025/2/14.
 */
public class DifyChatDialect extends AbstractChatDialect {

    @Override
    public boolean matched(ChatConfig config) {
        return "dify".equals(config.getProvider());
    }

    @Override
    public String buildRequestJson(ChatConfig config, ChatOptions options, List<ChatMessage> messages, boolean stream) {

        Optional<ChatMessage> first = messages.stream().limit(1).findFirst();

        if (!first.isPresent()) {
            return "";
        }

        return JSONUtil.createObj()
                .set("query", first.get().getContent())
                .set("response_mode", "streaming")
                .set("conversation_id", Optional.ofNullable(first.get().getMetadata().get("conversationId")).orElse(""))
                .set("user", Optional.ofNullable(first.get().getMetadata().get("user")).orElse("test"))
                .set("inputs", "")
                .toString();

    }

    @Override
    public boolean parseResponseJson(ChatConfig config, ChatResponseDefault resp, String respJson) {

        if (respJson.contains("data:")) {
            respJson = respJson.replaceFirst("data:", "");
            JSONObject parseObj = JSONUtil.parseObj(respJson);

            resp.setFinished("message_end".equals(parseObj.getStr("event")));

            AssistantMessage assistantMessage = new AssistantMessage(Optional.ofNullable(parseObj.getStr("answer")).orElse(""), false, null, null);

            if (resp.isFinished()) {

                assistantMessage.addMetadata("conversationId", parseObj.getStr("conversation_id"));

                if (parseObj.getJSONObject("metadata") != null) {
                    JSONArray files = parseObj.getJSONObject("metadata").getJSONArray("retriever_resources");
                    if (CollUtil.isNotEmpty(files)) {

                        String fileNames = files.stream().map(file -> JSONUtil.parseObj(file.toString()).getStr("document_name")).distinct().collect(Collectors.joining(","));

                        assistantMessage.addMetadata("file", fileNames);
                    }
                }

            }

            resp.addChoice(new ChatChoice(0, new Date(), parseObj.getStr("event"), assistantMessage));
        }

        return true;
    }


}
